module operators_mod

  use params_mod
  use recon_mod
  use math_mod
  implicit none

  type recon_matrics_type
    integer(i_kind), dimension(:,:,:), allocatable :: iCell           ! x index of reconstruction cells
    integer(i_kind), dimension(:,:,:), allocatable :: jCell           ! y index of reconstruction cells
    integer(i_kind), dimension(:,:  ), allocatable :: nRecCells       ! number of cells for reconstruction

    real   (r_kind), dimension(:,:,:,:), allocatable :: matrix

    integer(i_kind), dimension(:,:), allocatable :: xDir
    integer(i_kind), dimension(:,:), allocatable :: yDir

    integer(i_kind), dimension(:,:), allocatable :: reconType
  end type recon_matrics_type

  type(recon_matrics_type), dimension(:), allocatable :: weno_matrics

contains

  subroutine spatial_operator_init(mesh)

    type(cubed_sphere_mesh_type), intent(in) :: mesh

    integer(i_kind) dms, dme
    integer(i_kind) idom, idom_ngb
    
    logical :: is_full_stencil = .true.

    call mesh%get_params(dms=dms, dme=dme)

    allocate(weno_matrics (mesh%ndom))

    do idom = dms, dme
      if (mesh%domain(idom)%active) then
        call weno_matrics_init(mesh, mesh%domain(idom), weno_matrics(idom),is_full_stencil)
        if (mesh%domain(idom)%left_domain%ptr%initialized) then
          call weno_matrics_init(mesh, mesh%domain(idom)%left_domain%ptr, weno_matrics(mesh%domain(idom)%left_domain%ptr%id),is_full_stencil)
        end if
        if (mesh%domain(idom)%right_domain%ptr%initialized) then
          call weno_matrics_init(mesh, mesh%domain(idom)%right_domain%ptr, weno_matrics(mesh%domain(idom)%right_domain%ptr%id),is_full_stencil)
        end if
        if (mesh%domain(idom)%top_domain%ptr%initialized) then
          call weno_matrics_init(mesh, mesh%domain(idom)%top_domain%ptr, weno_matrics(mesh%domain(idom)%top_domain%ptr%id),is_full_stencil)
        end if
        if (mesh%domain(idom)%bottom_domain%ptr%initialized) then
          call weno_matrics_init(mesh, mesh%domain(idom)%bottom_domain%ptr, weno_matrics(mesh%domain(idom)%bottom_domain%ptr%id),is_full_stencil)
        end if
      end if
    end do
    
  end subroutine spatial_operator_init
  
  subroutine reconstruction_2d(weno_matrics,                 &
                               ims,ime,jms,jme,              &
                               ids,ide,jds,jde,              &
                               q,qL,qR,qB,qT,qQ,             &
                               dqdxL,dqdxR,dqdxB,dqdxT,dqdxQ,&
                               dqdyL,dqdyR,dqdyB,dqdyT,dqdyQ,&
                               dx, dy,                       &
                               trouble_cell)
    integer(i_kind)            , intent(in) :: ims,ime,jms,jme
    integer(i_kind)            , intent(in) :: ids,ide,jds,jde
    type   (recon_matrics_type), intent(in) :: weno_matrics
    real   (r_kind), dimension(                  ims:ime,jms:jme), intent(in   )          :: q
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: qL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: qR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: qB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: qT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme), intent(inout),optional :: qQ
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdxL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdxR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdxB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdxT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme), intent(  out),optional :: dqdxQ
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdyL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdyR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdyB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme), intent(  out),optional :: dqdyT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme), intent(  out),optional :: dqdyQ
    real   (r_kind)                                              , intent(in   ),optional :: dx
    real   (r_kind)                                              , intent(in   ),optional :: dy
    logical        , dimension(                  ims:ime,jms:jme), intent(in   ),optional :: trouble_cell

    real(r_kind), dimension(maxRecCells) :: u

    real   (r_kind), dimension(nWenoPoints) :: qrec
    real   (r_kind), dimension(nWenoPoints) :: qx
    real   (r_kind), dimension(nWenoPoints) :: qy
    real   (r_kind), dimension(maxRecCells) :: qH
    integer(i_kind) :: xdir,ydir
    logical         :: TC

    integer(i_kind) :: iVar,i,j,k,iCOS,iCell
    integer(i_kind) :: iRec,jRec
    
    if( present(dqdxL).and..not.present(dx) )then
      print*,'Present dqdxL but not present dx'
      stop
    endif
    
    if( present(dqdyL).and..not.present(dy) )then
      print*,'Present dqdyL but not present dy'
      stop
    endif

    !$OMP PARALLEL DO PRIVATE(i,iCOS,iRec,jRec,u,TC,iCell,qrec,qx,qy,qH) COLLAPSE(2)
    do j = jds,jde
      do i = ids,ide
        do iCOS = 1,maxRecCells
          iRec = weno_matrics%iCell(iCOS,i,j)
          jRec = weno_matrics%jCell(iCOS,i,j)
          u(iCOS) = q(iRec,jRec)
        enddo

        if(present(trouble_cell))then
          if( any( trouble_cell(weno_matrics%iCell(:,i,j),weno_matrics%jCell(:,i,j)) ) )then
            TC = .true.
          else
            TC = .false.
          endif
        else
          TC = .true.
        endif

        if(TC)then
          if( present(dqdxL) .or. present(dqdyL) )then
            call WENO(u,weno_matrics%reconType(i,j),qrec,qx,qy)
          else
            call WENO(u,weno_matrics%reconType(i,j),qrec)
          endif
        else
         iCOS = 0
         do iCell = 1,maxRecCells
           if(existPolyTerm(weno_matrics%reconType(i,j),iCell)==1)then
             iCOS = iCOS + 1
             qH(iCOS) = u(iCell)
           endif
         enddo

         qrec = matmul( iAPoly  (weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
         if(present(dqdxL))qx   = matmul( iAPoly_x(weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
         if(present(dqdyL))qy   = matmul( iAPoly_y(weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
        endif
        
        if(present(qL   )) qL   (:,i,j) = qrec(wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
        if(present(qR   )) qR   (:,i,j) = qrec(wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
        if(present(qB   )) qB   (:,i,j) = qrec(wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
        if(present(qT   )) qT   (:,i,j) = qrec(wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
        if(present(qQ   )) qQ   (:,i,j) = qrec(wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
        
        if(present(dqdxL)) dqdxL(:,i,j) = qx  (wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
        if(present(dqdxR)) dqdxR(:,i,j) = qx  (wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
        if(present(dqdxB)) dqdxB(:,i,j) = qx  (wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
        if(present(dqdxT)) dqdxT(:,i,j) = qx  (wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
        if(present(dqdxQ)) dqdxQ(:,i,j) = qx  (wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
        
        if(present(dqdyL)) dqdyL(:,i,j) = qy  (wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
        if(present(dqdyR)) dqdyR(:,i,j) = qy  (wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
        if(present(dqdyB)) dqdyB(:,i,j) = qy  (wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
        if(present(dqdyT)) dqdyT(:,i,j) = qy  (wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
        if(present(dqdyQ)) dqdyQ(:,i,j) = qy  (wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
      enddo
    enddo
    !$OMP END PARALLEL DO
  end subroutine reconstruction_2d
  
  subroutine reconstruction_h(weno_matrics,                  &
                              ims,ime,jms,jme,kms,kme,       &
                              ids,ide,jds,jde,kds,kde,       &
                              q,qL,qR,qB,qT,qQ,              &
                               dqdxL,dqdxR,dqdxB,dqdxT,dqdxQ,&
                               dqdyL,dqdyR,dqdyB,dqdyT,dqdyQ,&
                               dx, dy,                       &
                               trouble_cell)
    integer(i_kind)            , intent(in) :: ims,ime,jms,jme,kms,kme
    integer(i_kind)            , intent(in) :: ids,ide,jds,jde,kds,kde
    type   (recon_matrics_type), intent(in) :: weno_matrics
    real   (r_kind), dimension(                  ims:ime,jms:jme,kms:kme), intent(in   )          :: q
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: qL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: qR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: qB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: qT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme,kms:kme), intent(inout),optional :: qQ
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdxL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdxR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdxB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdxT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdxQ
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdyL
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdyR
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdyB
    real   (r_kind), dimension(nPointsOnEdge    ,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdyT
    real   (r_kind), dimension(nQuadPointsOnCell,ims:ime,jms:jme,kms:kme), intent(  out),optional :: dqdyQ
    real   (r_kind)                                                      , intent(in   ),optional :: dx
    real   (r_kind)                                                      , intent(in   ),optional :: dy
    logical        , dimension(                  ims:ime,jms:jme,kms:kme), intent(in   ),optional :: trouble_cell

    real(r_kind), dimension(maxRecCells) :: u

    real   (r_kind), dimension(nWenoPoints) :: qrec
    real   (r_kind), dimension(nWenoPoints) :: qx
    real   (r_kind), dimension(nWenoPoints) :: qy
    real   (r_kind), dimension(maxRecCells) :: qH
    integer(i_kind) :: xdir,ydir
    logical         :: TC

    integer(i_kind) :: iVar,i,j,k,iCOS,iCell
    integer(i_kind) :: iRec,jRec
    
    if( present(dqdxL).and..not.present(dx) )then
      print*,'Present dqdxL but not present dx'
      stop
    endif
    
    if( present(dqdyL).and..not.present(dy) )then
      print*,'Present dqdyL but not present dy'
      stop
    endif

    !$OMP PARALLEL DO PRIVATE(i,j,iCOS,iRec,jRec,u,TC,iCell,qrec,qx,qy,qH) COLLAPSE(3)
    do k = kds,kde
      do j = jds,jde
        do i = ids,ide
          do iCOS = 1,maxRecCells
            iRec = weno_matrics%iCell(iCOS,i,j)
            jRec = weno_matrics%jCell(iCOS,i,j)
            u(iCOS) = q(iRec,jRec,k)
          enddo

          if(present(trouble_cell))then
            if( any( trouble_cell(weno_matrics%iCell(:,i,j),weno_matrics%jCell(:,i,j),k) ) )then
              TC = .true.
            else
              TC = .false.
            endif
          else
            TC = .true.
          endif

          if(TC)then
            if( present(dqdxL) .or. present(dqdyL) )then
              call WENO(u,weno_matrics%reconType(i,j),qrec,qx,qy)
            else
              call WENO(u,weno_matrics%reconType(i,j),qrec)
            endif
          else
           iCOS = 0
           do iCell = 1,maxRecCells
             if(existPolyTerm(weno_matrics%reconType(i,j),iCell)==1)then
               iCOS = iCOS + 1
               qH(iCOS) = u(iCell)
             endif
           enddo

           qrec = matmul( iAPoly  (weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
           if(present(dqdxL))qx   = matmul( iAPoly_x(weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
           if(present(dqdyL))qy   = matmul( iAPoly_y(weno_matrics%reconType(i,j),:,1:iCOS), qH(1:iCOS) )
          endif

          if(present(qL   )) qL   (:,i,j,k) = qrec(wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
          if(present(qR   )) qR   (:,i,j,k) = qrec(wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
          if(present(qB   )) qB   (:,i,j,k) = qrec(wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
          if(present(qT   )) qT   (:,i,j,k) = qrec(wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
          if(present(qQ   )) qQ   (:,i,j,k) = qrec(wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j)))
          
          if(present(dqdxL)) dqdxL(:,i,j,k) = qx  (wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
          if(present(dqdxR)) dqdxR(:,i,j,k) = qx  (wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
          if(present(dqdxB)) dqdxB(:,i,j,k) = qx  (wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
          if(present(dqdxT)) dqdxT(:,i,j,k) = qx  (wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
          if(present(dqdxQ)) dqdxQ(:,i,j,k) = qx  (wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dx
          
          if(present(dqdyL)) dqdyL(:,i,j,k) = qy  (wenoLIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
          if(present(dqdyR)) dqdyR(:,i,j,k) = qy  (wenoRIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
          if(present(dqdyB)) dqdyB(:,i,j,k) = qy  (wenoBIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
          if(present(dqdyT)) dqdyT(:,i,j,k) = qy  (wenoTIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
          if(present(dqdyQ)) dqdyQ(:,i,j,k) = qy  (wenoQIdx(:,weno_matrics%xDir(i,j),weno_matrics%yDir(i,j))) / dy
        enddo
      enddo
    enddo
    !$OMP END PARALLEL DO
  end subroutine reconstruction_h

  subroutine weno_matrics_init(mesh, domain, weno_matrics,is_full_stencil)
    type(cubed_sphere_mesh_type  ), intent(in   ) :: mesh
    type(cubed_sphere_domain_type), intent(in   ) :: domain
    type(recon_matrics_type      ), intent(inout) :: weno_matrics
    logical, optional,              intent(in   ) :: is_full_stencil
    
    logical check_stencil
    integer(i_kind) ims, ime, jms, jme, kms, kme
    integer(i_kind) ids, ide, jds, jde, kds, kde
    integer(i_kind) irs, ire, jrs, jre
    integer(i_kind) dms, dme
    integer(i_kind) ipds, ipde, jpds, jpde
    integer(i_kind) idom

    integer(i_kind) i, j, k, iPOC, iCOS, iType, iRec, jRec, iR, jR, m, n, xdir, ydir

    integer(i_kind), dimension(:), allocatable :: lack

    check_stencil = .true.
    if(present(is_full_stencil))check_stencil = .not. is_full_stencil
    
    idom = domain%id

    call mesh%get_params(idom=idom                                  , &
                         ims=ims, ime=ime, jms=jms, jme=jme         , &
                         ids=ids, ide=ide, jds=jds, jde=jde         , &
                         kms=kms, kme=kme, kds=kds, kde=kde         , &
                         irs=irs, ire=ire, jrs=jrs, jre=jre         , &
                         dms=dms, dme=dme                           , &
                         ipds=ipds, ipde=ipde, jpds=jpds, jpde=jpde)

    if (.not. domain%touch_left_edge  ) ids = ids - 1
    if (.not. domain%touch_right_edge ) ide = ide + 1
    if (.not. domain%touch_top_edge   ) jde = jde + 1
    if (.not. domain%touch_bottom_edge) jds = jds - 1
    
    ! Just for geog interpolation
    if(.not.check_stencil)then
      ids = irs
      ide = ire
      jds = jrs
      jde = jre
    endif

    allocate(lack(nWenoLack))
    if(.not.allocated(weno_matrics%iCell    ))allocate(weno_matrics%iCell    (maxRecCells,ids:ide,jds:jde))
    if(.not.allocated(weno_matrics%jCell    ))allocate(weno_matrics%jCell    (maxRecCells,ids:ide,jds:jde))
    if(.not.allocated(weno_matrics%xDir     ))allocate(weno_matrics%xDir     (            ids:ide,jds:jde))
    if(.not.allocated(weno_matrics%yDir     ))allocate(weno_matrics%yDir     (            ids:ide,jds:jde))
    if(.not.allocated(weno_matrics%reconType))allocate(weno_matrics%reconType(            ids:ide,jds:jde))

    do j = jds, jde
      do i = ids, ide
        iCOS = 0
        iPOC = 0
        weno_matrics%reconType(i,j) = 1

        weno_matrics%xDir(i,j) = 1
        weno_matrics%yDir(i,j) = 1
        
        if(check_stencil)then
          if( i<ipds+recBdy.and.j<jpds+recBdy )then ! low left corner
            weno_matrics%xDir(i,j) = -1
            weno_matrics%yDir(i,j) = -1
          elseif(i>ipde-recBdy.and.j<jpds+recBdy)then ! low right corner
            weno_matrics%xDir(i,j) = 1
            weno_matrics%yDir(i,j) = -1
          elseif(i>ipde-recBdy.and.j>jpde-recBdy)then ! up right corner
            weno_matrics%xDir(i,j) = 1
            weno_matrics%yDir(i,j) = 1
          elseif(i<ipds+recBdy.and.j>jpde-recBdy)then ! up left corner
            weno_matrics%xDir(i,j) = -1
            weno_matrics%yDir(i,j) = 1
          endif
          
          xdir = weno_matrics%xDir(i,j)
          ydir = weno_matrics%yDir(i,j)
          
          lack = 0
          do jRec = -recBdy,recBdy
            do iRec = -recBdy,recBdy
              iCOS = iCOS + 1
              iR = i + xdir * iRec
              jR = j + ydir * jRec
          
              if( domain%cell_type(iR,jR) == corner_cell )then
                iPOC = iPOC + 1
                lack(iPOC) = iCOS
              endif
          
              weno_matrics%iCell(iCOS,i,j) = iR
              weno_matrics%jCell(iCOS,i,j) = jR
            enddo
          enddo
          
          do iType = 2,nWenoType
            if( sum( lack - wenoLack(iType,:) )==0 )then
              weno_matrics%reconType(i,j) = iType
              exit
            endif
          enddo
          
        else
          xdir = weno_matrics%xDir(i,j)
          ydir = weno_matrics%yDir(i,j)
          
          do jRec = -recBdy,recBdy
            do iRec = -recBdy,recBdy
              iCOS = iCOS + 1
              iR = i + xdir * iRec
              jR = j + ydir * jRec
          
              weno_matrics%iCell(iCOS,i,j) = iR
              weno_matrics%jCell(iCOS,i,j) = jR
            enddo
          enddo
          
        endif
        
      enddo
    enddo

  end subroutine weno_matrics_init

end module operators_mod
